package com.zyt.wiki.interceptor.rate;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 令牌桶对象（需要限流的每个Url或IP，都会对应一个桶对象）
 */
public class TokenBucket {

    private long bucketLimit ;	// 桶最大容量
    private long tokenPerInterval; // 单次放入令牌数量
    private long intervalInMills ; 	// 间隔时间（单位：毫秒）
    private long lastRefillTime ; 	// 最近填充时间
    private long leftToken ; 		// 剩余令牌数
    private String bucketType ; // 令牌桶类型名：token-ip:IP令牌桶，token-url：URL令牌桶

    private Logger logger = LoggerFactory.getLogger(this.getClass()) ;

    /**
     * @param bucketType	令牌桶类型名
     * @param bucketLimit	桶容量
     * @param tokenPerInterval	单次放入令牌数量
     * @param intervalInMills	间隔时间
     */
    public TokenBucket(String bucketType ,long bucketLimit, long tokenPerInterval, long intervalInMills) {
        this.bucketType = bucketType ;
        this.bucketLimit = bucketLimit ;
        this.tokenPerInterval = tokenPerInterval ;
        this.intervalInMills = intervalInMills ;

        // 首次默认桶容量是满的，达到预热效果
        this.lastRefillTime = System.currentTimeMillis() ;
        this.leftToken = bucketLimit ;
    }

    /**
     *  申请令牌
     *  （同一个IP或URL，对应的是同一个TokenBucket对象，所以要进行同步控制）
     * @return  true：有令牌    false：没有令牌
     */
    public synchronized boolean tryAcquire() {
        long now = System.currentTimeMillis() ; // 获取当前时间
        long generatedToken = (now-this.lastRefillTime)/this.intervalInMills * this.tokenPerInterval ; // 最近请求间隔 * 放入令牌速率
        this.leftToken = Math.min(this.bucketLimit, this.leftToken + generatedToken) ;
        this.logger.info("【RateType: {}】 leftToken: {}", this.bucketType, this.leftToken);
        if(leftToken >= 1) {
            this.lastRefillTime = now ;
            this.leftToken -- ;
            return true ;
        }else {
            return false ;
        }
    }

}
